﻿using System.CodeDom.Compiler;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using Unity.Entities.SourceGen.Common;
using Unity.Entities.SourceGen.SystemGenerator.Common;
using static Unity.Entities.SourceGen.JobEntityGenerator.JobEntityModule;

namespace Unity.Entities.SourceGen.JobEntityGenerator;

/*
 The `IjeSchedulingWalker` traverses through syntax nodes that have been marked by the `JobEntityModule` as candidates for patching -- specifically,
 scheduling invocations of `IJobEntity` instances. The `IjeSchedulingWalker`, like the `SystemApiWalker`, needs to handle nested candidates, e.g.

     new UsingCodeGenInsideAScheduleObjectInitJob
    {
        datas = SystemAPI.GetComponentLookup<EcsTestData>()
    }.Schedule(SystemAPI.QueryBuilder().WithAll<EcsTestData>().Build(), Dependency).Complete();

 In order to patch the example above correctly, the `IjeSchedulingSyntaxWalker` needs to cede write control to the `SystemApiContextWalker` when it reaches
 the `SystemAPI.GetComponentLookup<EcsTestData>()` node, and to the `SystemApiQueryBuilderSyntaxWalker` when it reaches the `SystemAPI.QueryBuilder().WithAll<EcsTestData>().Build()`
 node. It then gathers the patches made by the two other walkers, and ends up writing the following code to replace the original code above:

     __ScheduleViaJobChunkExtension_15(new UsingCodeGenInsideAScheduleObjectInitJob
        {
            datas = global::Unity.Entities.Internal.InternalCompilerInterface.GetComponentLookup<global::Unity.Entities.Tests.EcsTestData>(ref __TypeHandle.__Unity_Entities_Tests_EcsTestData_RW_ComponentLookup, ref this.CheckedStateRef)
        },
        __query_1239178100_0, Dependency, ref this.CheckedStateRef, true).Complete()

 The resultant `global::Unity.Entities.Internal.InternalCompilerInterface.GetComponentLookup<global::Unity.Entities.Tests.EcsTestData>(ref __TypeHandle.__Unity_Entities_Tests_EcsTestData_RW_ComponentLookup, ref this.CheckedStateRef)`
 was generated by the `SystemApiContextWalker`, and `__query_1239178100_0` by the `SystemApiQueryBuilderSyntaxWalker`.
 */
public class IjeSchedulingSyntaxWalker : CSharpSyntaxWalker, IModuleSyntaxWalker
{
    private readonly Dictionary<InvocationExpressionSyntax, JobEntityInstanceInfo> _schedulingInvocationNodes;
    private readonly StringWriter _schedulingArgsInnerWriter;
    private readonly IndentedTextWriter _schedulingArgsWriter;

    private IndentedTextWriter _writer;
    private bool _hasWrittenSyntax;
    private int _uniqueId;
    private bool _isWalkingSchedulingInvocationArgument;
    private SystemDescription _systemDescription;

    private string _schedulingJobEntityInstanceArgument;
    private string _userDefinedQueryArgument;
    private string _userDefinedDependency;
    private ObjectCreationExpressionSyntax _jobEntityInstanceCreationSyntax;

    internal IjeSchedulingSyntaxWalker(ref SystemDescription systemDescription, Dictionary<InvocationExpressionSyntax, JobEntityInstanceInfo> jobEntityInfos)
        : base(SyntaxWalkerDepth.Trivia)
    {
        _uniqueId = 0;
        _systemDescription = systemDescription;
        _schedulingInvocationNodes = jobEntityInfos;

        _schedulingArgsInnerWriter = new StringWriter();
        _schedulingArgsWriter = new IndentedTextWriter(_schedulingArgsInnerWriter);
    }

    public bool TryWriteSyntax(IndentedTextWriter writer, CandidateSyntax candidateSyntax)
    {
        _writer = writer;
        _hasWrittenSyntax = false;

        _schedulingArgsInnerWriter.GetStringBuilder().Clear();
        _isWalkingSchedulingInvocationArgument = false;

        // Begin depth-first traversal of the candidate node
        Visit(candidateSyntax.Node);

        return _hasWrittenSyntax;
    }

    public override void VisitInvocationExpression(InvocationExpressionSyntax node)
    {
        if (_isWalkingSchedulingInvocationArgument)
        {
            if (_systemDescription.CandidateNodes.TryGetValue(node, out var candidateSyntax))
            {
                var rewroteArg = _systemDescription.SyntaxWalkers[candidateSyntax.GetOwningModule()].TryWriteSyntax(_schedulingArgsWriter, candidateSyntax);
                if (!rewroteArg)
                    base.VisitInvocationExpression(node);
            }
            else
                base.VisitInvocationExpression(node);
        }
        else if (_schedulingInvocationNodes.TryGetValue(node, out var jobEntityInfo))
        {
            var tryGetJobArg = jobEntityInfo.TryGetJobArgumentUsedInSchedulingInvocation();
            if (tryGetJobArg.Success)
            {
                if (tryGetJobArg.ObjectCreationExpressionSyntax != null)
                {
                    _isWalkingSchedulingInvocationArgument = true;

                    /* The user might have used sourcegen-dependent API when creating a new `IJobEntity` instance, e.g.

                        new UsingCodeGenInsideAScheduleObjectInitJob
                        {
                            datas = SystemAPI.GetComponentLookup<EcsTestData>()
                        }.Schedule(SystemAPI.QueryBuilder().WithAll<EcsTestData>().Build(), Dependency).Complete();

                    Thus we need to walk the object creation syntax, and cede write control to the appropriate syntax walker if necessary.
                    In the example above, we need to cede write control to the `SystemApiContextWalker` when we reach the
                    `SystemAPI.GetComponentLookup<EcsTestData>()` node, so that the `SystemApiContextWalker` can patch it accordingly.
                    */
                    _jobEntityInstanceCreationSyntax = tryGetJobArg.ObjectCreationExpressionSyntax;
                    VisitObjectCreationExpression(_jobEntityInstanceCreationSyntax);

                    _isWalkingSchedulingInvocationArgument = false;
                }
                else
                {
                    _jobEntityInstanceCreationSyntax = null;
                    _schedulingJobEntityInstanceArgument = tryGetJobArg.IdentifierNameSyntax.ToString();
                }
            }

            var result =
                JobEntityInstanceInfo.GetUserDefinedQueryAndDependency(ref _systemDescription, jobEntityInfo.IsExtensionMethodUsed, node);

            if (result.UserDefinedEntityQuery != null)
            {
                _isWalkingSchedulingInvocationArgument = true;

                // The user might have used sourcegen-dependent API to create a query, so we need to walk it,
                // and cede write control to the appropriate syntax walker if necessary. In the example shown in the
                // previous comment, we will need to cede write control to the `SystemApiQueryBuilderWalker` when
                // we reach the `SystemAPI.QueryBuilder().WithAll<EcsTestData>().Build()` node, so that the
                // `SystemApiQueryBuilderWalker` can patch it accordingly.
                base.Visit(result.UserDefinedEntityQuery);

                _userDefinedQueryArgument = _schedulingArgsInnerWriter.ToString();

                // StringWriter.Flush() unfortunately doesn't clear the buffer correctly: https://stackoverflow.com/a/13706647
                // Since IndentedTextWriter.Flush() calls stringWriter.Flush(), we cannot use it either.
                _schedulingArgsInnerWriter.GetStringBuilder().Clear();
                _isWalkingSchedulingInvocationArgument = false;
            }
            else
                _userDefinedQueryArgument = null;

            _userDefinedDependency = result.UserDefinedDependency?.ToString();

            string replacementCode =
                jobEntityInfo.GetAndAddScheduleExpression(
                    ref _systemDescription,
                    _uniqueId++,
                    tryGetJobArg.JobArgumentIndexInExtensionMethod,
                    _schedulingJobEntityInstanceArgument,
                    _userDefinedQueryArgument,
                    _userDefinedDependency);

            _writer.Write(replacementCode);
            _hasWrittenSyntax = true;
        }
        else
            _hasWrittenSyntax = false;
    }

    public override void VisitObjectCreationExpression(ObjectCreationExpressionSyntax node)
    {
        base.VisitObjectCreationExpression(node);

        if (node == _jobEntityInstanceCreationSyntax)
        {
            _schedulingJobEntityInstanceArgument = _schedulingArgsInnerWriter.ToString();

            // StringWriter.Flush() unfortunately doesn't clear the buffer correctly: https://stackoverflow.com/a/13706647
            // Since IndentedTextWriter.Flush() calls stringWriter.Flush(), we cannot use it either.
            _schedulingArgsInnerWriter.GetStringBuilder().Clear();
        }
    }

    public override void VisitMemberAccessExpression(MemberAccessExpressionSyntax node)
    {
        if (_isWalkingSchedulingInvocationArgument)
        {
            if (_systemDescription.CandidateNodes.TryGetValue(node, out var candidateSyntax))
            {
                var rewroteArg = _systemDescription.SyntaxWalkers[candidateSyntax.GetOwningModule()].TryWriteSyntax(_schedulingArgsWriter, candidateSyntax);
                if (!rewroteArg)
                    base.VisitMemberAccessExpression(node);
            }
            else
                base.VisitMemberAccessExpression(node);
        }
        else
            base.VisitMemberAccessExpression(node);
    }

    public override void VisitIdentifierName(IdentifierNameSyntax node)
    {
        if (_isWalkingSchedulingInvocationArgument)
        {
            if (_systemDescription.CandidateNodes.TryGetValue(node, out var candidateSyntax))
            {
                var rewroteArg = _systemDescription.SyntaxWalkers[candidateSyntax.GetOwningModule()].TryWriteSyntax(_schedulingArgsWriter, candidateSyntax);
                if (!rewroteArg)
                    base.VisitIdentifierName(node);
            }
            else
                base.VisitIdentifierName(node);
        }
        else
            base.VisitIdentifierName(node);
    }
    public override void VisitToken(SyntaxToken token)
    {
        VisitLeadingTrivia(token);

        if (_isWalkingSchedulingInvocationArgument)
            _schedulingArgsWriter.Write(token.Text);
        else
            _writer.Write(token.Text);

        VisitTrailingTrivia(token);
    }

    public override void VisitTrivia(SyntaxTrivia trivia)
    {
        var triviaKind = trivia.Kind();

        if (triviaKind == SyntaxKind.EndOfLineTrivia)
        {
            if (_isWalkingSchedulingInvocationArgument)
                _schedulingArgsWriter.WriteLine();
            else
                _writer.WriteLine();
        }

        else if (triviaKind != SyntaxKind.DisabledTextTrivia &&
                 triviaKind != SyntaxKind.PreprocessingMessageTrivia &&
                 triviaKind != SyntaxKind.IfDirectiveTrivia &&
                 triviaKind != SyntaxKind.ElifDirectiveTrivia &&
                 triviaKind != SyntaxKind.ElseDirectiveTrivia &&
                 triviaKind != SyntaxKind.EndIfDirectiveTrivia &&
                 triviaKind != SyntaxKind.RegionDirectiveTrivia &&
                 triviaKind != SyntaxKind.EndRegionDirectiveTrivia &&
                 triviaKind != SyntaxKind.DefineDirectiveTrivia &&
                 triviaKind != SyntaxKind.UndefDirectiveTrivia &&
                 triviaKind != SyntaxKind.ErrorDirectiveTrivia &&
                 triviaKind != SyntaxKind.WarningDirectiveTrivia &&
                 triviaKind != SyntaxKind.PragmaWarningDirectiveTrivia &&
                 triviaKind != SyntaxKind.PragmaChecksumDirectiveTrivia &&
                 triviaKind != SyntaxKind.ReferenceDirectiveTrivia &&
                 triviaKind != SyntaxKind.BadDirectiveTrivia &&
                 triviaKind != SyntaxKind.SingleLineCommentTrivia &&
                 triviaKind != SyntaxKind.MultiLineCommentTrivia)
        {
            if (!trivia.HasStructure)
            {
                if (_isWalkingSchedulingInvocationArgument)
                    _schedulingArgsWriter.Write(trivia.ToString());
                else
                    _writer.Write(trivia.ToString());
            }
        }
    }
}
